#!/usr/bin/env python
"""
        AlignedSent - 
        This class is responsible for storing and handling the translation 
        alignments between two sentences 
		
		Baolin Jing, Kai Chen

"""

try:
    import pylab
    __pylab_loaded__ = True
except ImportError:
    __pylab_loaded__ = False

class Alignment(dict):
    """ the alignment class
    """
    
    def __init__(self, align):
        """ initializer       
        # e.g. for non-empty alignment
        # Alignment({(0,0):True, (1,0):False, (2,1):True, (3,2):True })
        #
        # where True is a sure alignment, vice versa
        """
        self = dict.__init__(self, align)

    def __repr__(self):
        return 'Alignment(' + dict.__repr__(self) + ')'

    def invert(self):
        d = {}
        for i,j in self:
            d[(j,i)] = self[(i,j)]
        return Alignment(d)


        
class AlignedSent(object):
    """
        The sentence alignment class.
        http://code.google.com/p/clir/wiki/AlignedSent?ts=1272297554&updated=AlignedSent
    """


    # initializer
    def __init__(self, e = [], f = [], align = Alignment({})):
        """Initializer.
        e = list of english words
        f = list of foreign words
        align = set of alignments
        
        # e.g.
        #
        # text_e = ['a','b','c','d']
        # text_f = ['1','2','3']
        # align = {(0,0):True, (1,0):False, (2,1):True, (3,2):True ])
        #
        # the True/False above stand for 'sure' or 'probable' alignments
        #
        """
        self.text_e = e
        self.text_f = f
        self.align = align

        self._check_align(align)
        
    def _check_align(self, a):
        # check legality of alignments
        for (i,j) in a:
            if i > len(self.text_e) - 1 or j > len(self.text_f) - 1:
                raise IndexError("Alignment out of sentence boundary")


    def __repr__(self):
        """s = '' + repr(self.text_e) + '\n' + repr(self.text_f) + '\n' \
            + repr(self.align)
        """
        s = 'AlignedSent(\n' +\
            repr(self.text_e) + ',\n' + \
            repr(self.text_f) + ',\n' + \
            repr(self.align) + '\n)'
        return s

    def gui_repr(self):
        """Generate a GUI to represent the sentence alignments
        """
        if __pylab_loaded__:
            fig_width = max(len(self.text_e), len(self.text_f)) + 1
            fig_height = 3
            pylab.figure(figsize=(fig_width*0.8, fig_height*0.8), facecolor='w')
            pylab.box(on=False)
            pylab.subplots_adjust(left=0, right=1, bottom=0, top=1)
            pylab.xlim(-1, fig_width - 1)
            pylab.ylim(0, fig_height)
            pylab.xticks([])
            pylab.yticks([])

            e = [0 for _ in xrange(len(self.text_e))]
            f = [0 for _ in xrange(len(self.text_f))]
            for (i, j) in self.align:
                e[i] = 1
                f[j] = 1
                # draw the middle line
                pylab.arrow(i, 2, j - i, -1, color='r')
            for i in xrange(len(e)):
                # draw e side line
                pylab.text(i, 2.5, self.text_e[i], ha = 'center', va = 'center',
                        rotation=30)
                if e[i] == 1:
                    pylab.arrow(i, 2.5, 0, -0.5, color='r', alpha=0.3, lw=2)
            for i in xrange(len(f)):
                # draw f side line
                 pylab.text(i, 0.5, self.text_f[i], ha = 'center', va = 'center',
                        rotation=30)
                 if f[i] == 1:
                    pylab.arrow(i, 0.5, 0, 0.5, color='r', alpha=0.3, lw=2)

            pylab.draw()

    @staticmethod
    def block_and_view_gui_repr():
        """Block the program to view previous spawned Repr GUIs
        """
        if __pylab_loaded__:
            pylab.show()

    def update_align(self, e,f,s = True):
        """to add a new alignment or change the params.
        s (sure/probable) is True by default.
        """
        self._check_align(Alignment({(e,f):s}))
        self.align.update(Alignment({(e,f):s}))
        
    def remove_align(self, e,f):
        """remove target Alignment with key (e,f).
        return the value of the alignment if exists.
        """
        try:
            return self.align.pop((e,f))
        except KeyError:
            return None

    def clear_align(self):
        """remove all alignment pairs
        """
        self.align = Alignment({})
        

    def align_pairs(self):
        """return a pure set of alignments in the sentence.
        """
        return set(self.align.keys())
 
    ### the converse factory functions
    def inv_AlignSent(self):
        """return a conversed AlignedSent with
        English and foreign language swapped.
        """
        return AlignedSent(self.text_f,self.text_e,self.align.invert())

    def inv_aligns(self):
        """ just the swapped alignment """
        return self.align.invert()


from collections import defaultdict

class AlignedSentCollection(object):


    def __init__(self):
        """ initializer
        initialize as empty collection
        
        #
        # e.g. of a non-empty AlignedSentCollection
        #
        # set_words_e = set(['oh','my','god'])
        # set_words_f = set(['ma','mi','a'])
        #
        # self.sents = defaultdict(<class 'AlignedSent'>, {0: Sent0, 1: Sent1})
        # where Sent0 is the first AlignedSent object, Sent1 is the second
        #

        """
        self.set_words_e = set([])
        self.set_words_f = set([])
        self.sents = defaultdict(AlignedSent)

    def __repr__(self):
        return str(self.sents)

    def refresh_vocabulary(self):
        e = set([])
        f = set([])
        for x in self.sents.values():
            for i in x.text_e:
                e.update([i])
            for j in x.text_f:
                f.update([j])
        self.set_words_e = e
        self.set_words_f = f

    def add(self, sent_num, aligned_sent):
        """
        add a new AlignedSent instance into collection
        """
        if sent_num not in self.sents:
            self.sents.update({sent_num:aligned_sent})
            for i in aligned_sent.text_e:
                self.set_words_e.update(i)
            for j in aligned_sent.text_f:
                self.set_words_f.update(j)
        else:
            raise KeyError("Target Sentence already exist, use update() to update")

    def remove(self, sent_num):
        """ remove is very costy due to vocabulary maintainance,
        try not to call remove() frequently.
        """
        try:
            self.sents.pop(sent_num)
            self.refresh_vocabulary()
        except KeyError:
            return None

    def update(self, sent_num, aligned_sent):
        """ update is very costy due to vocabulary maintainance,
        try not to call update() frequently.
        """
        self.remove(sent_num)
        self.add(sent_num, aligned_sent)
                

def naacl_collection_parser(naacl_e, naacl_f, naacl_wa):
    collection = AlignedSentCollection()

    for i in range(0, len(naacl_e.sents())):
        sent_wa = naacl_wa.sents()[i][1:]
        align = Alignment({})
        for j in range(1, len(sent_wa),3):
            align.update({(int(sent_wa[j-1]) - 1, int(sent_wa[j]) - 1): \
                          sent_wa[j+1] == 'S'})
        sent_e = naacl_e.sents()[i]
        sent_f = naacl_f.sents()[i]
        sent_number = sent_e.pop()
        sent_f.pop()
        aligned_sent = AlignedSent(sent_e, sent_f, align)
        collection.add(sent_number, aligned_sent)
    return collection
    

if __name__ == "__main__":
    print "=================\nclass Alignment:\n================="
    a = Alignment({(1,2):True, (2,3):False})
    print a
    print "a.invert()\n",a.invert(),'\n'
    print "=================\nclass AlignedSent:\n================="
    s = AlignedSent(['Joan','of','arts'],['Joan','de','arc'],\
                    Alignment({(0,0):True, (1,1):True, (2,2):False}))
    print s
    s.update_align(2,2)
    print "updated the 3rd alignment from Prabable to Sure"
    print s,'\n'

    print "=================\nclass AlignedSentCollection:\n================="


    from aligned_corpus_reader import *

    root = '../../data/naacl/'
    align_fileid = ['trial.wa']
    english_fileids = ['trial.e']
    foreign_fileids = ['trial.f']
    naacl_mode = 'text'
    naacl_mode_1 = 'align'
    naacl_e = NaaclCorpusReader(root, english_fileids, naacl_mode)
    naacl_f = NaaclCorpusReader(root, foreign_fileids, naacl_mode)
    naacl_wa = NaaclCorpusReader(root, align_fileid, naacl_mode_1)
##    print naacl_wa.sents()
    collection = naacl_collection_parser(naacl_e, naacl_f, naacl_wa)
    collection.refresh_vocabulary()
##    print collection.set_words_e
    collection.sents['1001'].gui_repr()
    AlignedSent.block_and_view_gui_repr()
#    naacl = NaaclCorpusReader(root, text_fileids, naacl_mode)
#    naacl_1 = NaaclCorpusReader(root, align_fileid, naacl_mode_1)
##    print naacl.words()
##    print len(naacl.words())
#    print naacl.sents()[0]
#    print naacl_1.sents()[-1]
##    for word in naacl.words(): print word
